package lib

func init() {
	Probs = append(Probs, Problem{
		Num:         327,
		Discription: "子数组之和在指定范围的个数",
		Level:       4,
		Labels: map[string]int{
			"前缀和":  1,
			"归并排序": 1,
		},
	})
}

// 问题可以转化为：有多少对（i,j），j>i, 满足preSum[j]-preSum[i]在指定范围内
// 在对preSum进行归并排序时进行计数即可

/*
计数：
在进行已经升序的左半边数组和已经升序的右半边数组合并时，计数方法如下：
依次遍历左半边的元素：
在确定i的情况下，求[left, right)为右半边数组满足条件的范围
初始left，right指向右半边数组的第一个元素
通过left++找到第一个>=lower的j（left）
再从left的位置开始，right++，找到第一个不满足<=upper的j（right）
那么在确定i的情况下，j有效的范围就是[left,right)，组合总数为right-left
在下一轮遍历中，i+1，此时不需要重置left，因为PreSum[i+1]>preSum[i],left前面的元素时是不可能满足preSum[left]-preSum[i+1]>=lower的（因为已有preSum[left-1]-preSum[i]<lower，必然有preSum[left-1]-preSum[i+1]<lower）
*/

/*
为什么在合并时只需要对左半边的元素进行计数：
若i j都在同一个半区，统计出来的对子上一轮归并有统计过
因此i必须在左半区，j必须在右半区
*/
func CountRangeSum(nums []int, lower int, upper int) int {
	count := 0

	var mergeSort func([]int)
	mergeSort = func(array []int) {
		n := len(array)
		if n <= 1 {
			return
		}

		mid := n / 2
		array1 := make([]int, mid)
		array2 := make([]int, n-mid)
		copy(array1, array[:mid])
		copy(array2, array[mid:])
		mergeSort(array1)
		mergeSort(array2)

		left := 0
		right := 0
		//left 不需要重置
		for i := 0; i < mid; i++ {
			for left < n-mid && array2[left]-array1[i] < lower {
				left++
			}
			//从left的位置继续找right
			right = left
			for right < n-mid && array2[right]-array1[i] <= upper {
				right++
			}
			count += right - left
		}

		i := 0
		j := 0
		k := 0
		for k < n {
			if i == mid {
				array[k] = array2[j]
				j++
				k++
				continue
			}

			if j == n-mid {
				array[k] = array1[i]
				i++
				k++
				continue
			}

			if array1[i] < array2[j] {
				array[k] = array1[i]
				i++
			} else {
				array[k] = array2[j]
				j++
			}
			k++
		}
	}

	//一定得len(nums)+1
	//为什么需要一个preSum[0]=0:考虑到有从0开始的子数组（不需要减去任何前缀和）的结果
	preSum := make([]int, len(nums)+1)
	preSum[0] = nums[0]
	for i := 1; i <= len(nums); i++ {
		preSum[i] = preSum[i-1] + nums[i-1]
	}
	mergeSort(preSum)
	return count
}
